import os
import pandas as pd
import matplotlib.pyplot as plt
import datetime
import numpy as np
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from utils.log import Logger
from utils.common import data_preprocessing
import xgboost as xgb
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score, GridSearchCV
from sklearn.metrics import roc_auc_score, accuracy_score, recall_score, f1_score, classification_report, roc_curve
import joblib
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15
from lightgbm import LGBMClassifier
from sklearn.tree import DecisionTreeClassifier

class PowerLoadModel(object):
    def __init__(self, path):
        logfile_name = 'train_' + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logfile = Logger('../', logfile_name).get_logger()
        self.data_source = data_preprocessing(path)

def feature_engineering(data):
    """
    核心特征：BusinessTravel、OverTime、MaritalStatus
    补充特征：Department、JobLevel、MonthlyIncome
    可尝试特征交互：
    OverTime × MonthlyIncome（加班与收入的平衡）
    JobLevel × TotalWorkingYears（职业发展匹配度）
    :param data:
    :param logger:
    :return:
    """
    feature_data = data.copy()
    # 商务旅行情况
    travel_dummies = pd.get_dummies(feature_data['BusinessTravel'], prefix='Travel')
    feature_data = pd.concat([feature_data, travel_dummies], axis=1)
    # 是否加班
    feature_data['OverTime'] = feature_data['OverTime'].map({'Yes': 1, 'No': 0})
    # 性别
    feature_data['Gender'] = feature_data['Gender'].map({'Male': 1, 'Female': 0})
    # 婚姻状况
    marital_dummies = pd.get_dummies(feature_data['MaritalStatus'], prefix='Marital')
    feature_data = pd.concat([feature_data, marital_dummies], axis=1)
    # 部门
    department_dummies = pd.get_dummies(feature_data['Department'], prefix='Dept')
    feature_data = pd.concat([feature_data, department_dummies], axis=1)
    # 工作等级
    job_level_dummies = pd.get_dummies(feature_data['JobLevel'], prefix='JobLevel')
    feature_data = pd.concat([feature_data, job_level_dummies], axis=1)
    # EducationField
    education_field_dummies = pd.get_dummies(feature_data['EducationField'], prefix='education')
    feature_data = pd.concat([feature_data, education_field_dummies], axis=1)

    # 添加:
    # 添加工作角色特征
    jobrole_dummies=pd.get_dummies(feature_data['JobRole'], prefix='jobrole')
    feature_data = pd.concat([feature_data, jobrole_dummies], axis=1)

    # OverTime × MonthlyIncome（加班与收入的平衡）
    feature_data['OvertimeIncome'] = feature_data['OverTime'] * feature_data['MonthlyIncome']
    # JobLevel × TotalWorkingYears（职业发展匹配度）
    feature_data['CareerMatch'] = feature_data['JobLevel'] * feature_data['TotalWorkingYears']

    # 构建特征列列表

    # 原始:
    # feature_columns = ['Attrition', 'Age', 'OverTime', 'DistanceFromHome',
    #                    'CareerMatch', 'EnvironmentSatisfaction', 'Gender', 'JobInvolvement', 'JobSatisfaction',
    #                    'NumCompaniesWorked', 'PerformanceRating', 'RelationshipSatisfaction', 'StockOptionLevel',
    #                    'TotalWorkingYears', 'TrainingTimesLastYear',
    #                    'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole', 'YearsSinceLastPromotion',
    #                    'YearsWithCurrManager'] + list(travel_dummies.columns) + list(marital_dummies.columns) + list(
    #     department_dummies.columns) + list(job_level_dummies.columns) + list(education_field_dummies.columns)

    feature_columns = (['Attrition',
                       'Age',
                       'OverTime',
                       # 'DistanceFromHome',
                       # 'CareerMatch',
                       'EnvironmentSatisfaction',
                       'Gender',
                        'JobInvolvement',
                       'JobSatisfaction',
                       'NumCompaniesWorked',
                       'PerformanceRating',
                       'RelationshipSatisfaction',
                       'StockOptionLevel',
                       # 'TotalWorkingYears',
                       # 'TrainingTimesLastYear',
                       'WorkLifeBalance',
                       # 'YearsAtCompany',
                       'YearsInCurrentRole',
                       # 'YearsSinceLastPromotion',
                       # 'YearsWithCurrManager',
                        'Education',
                        'PercentSalaryHike']
                        # 'StandardHours'] # 有没有无影响

                       + list(travel_dummies.columns)  # BusinessTravel
                       + list(marital_dummies.columns) # MaritalStatus
                       + list(department_dummies.columns)  # Department
                       + list(job_level_dummies.columns)  # MaritalStatus
                       + list(education_field_dummies.columns) # EducationField
                       + list(jobrole_dummies.columns)) # JobRole

    return feature_data[feature_columns], feature_columns
def model_test(test_path, features):
    # 加载测试数据
    test_data = data_preprocessing(test_path)
    # 特征工程
    test_feature_data, _ = feature_engineering(test_data)
    x_test = test_feature_data[features[1:]]
    y_test = test_feature_data[features[0]]

    # 加载模型
    model = joblib.load('../model/rc_20250607.pkl')
    # model1 = joblib.load('../model/rc_20250607_gbdt.pkl')

    # 预测概率
    y_pred_proba = model.predict_proba(x_test)[:, 1]
    y_pred = model.predict(x_test)

    # 模型评估
    auc_score = roc_auc_score(y_test, y_pred_proba)

    print(f"Test AUC Score: {auc_score}")
    # print(classification_report(y_test, y_pred))
    # 计算ROC曲线
    fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)

    # 绘制ROC曲线
    plt.figure(figsize=(10, 8))  # 增大图表尺寸
    plt.plot(fpr, tpr, color='darkorange', lw=2.5,
             label=f'ROC曲线 (AUC = {auc_score:.3f})')  # 使用更鲜明的颜色和更粗的线条
    plt.plot([0, 1], [0, 1], color='navy', lw=1.5, linestyle='--',
             label='随机猜测线 (AUC = 0.5)')  # 美化参考线

    # 添加最佳阈值点
    optimal_idx = np.argmax(tpr - fpr)
    optimal_threshold = thresholds[optimal_idx]
    plt.plot(fpr[optimal_idx], tpr[optimal_idx], 'go', markersize=10,
             label=f'最佳阈值点: {optimal_threshold:.3f}')  # 标记最佳阈值

    # 添加网格线
    plt.grid(True, linestyle='--', alpha=0.6)

    # 调整坐标轴范围和标签
    plt.xlim([-0.01, 1.01])  # 略微扩展坐标轴范围
    plt.ylim([-0.01, 1.01])
    plt.xlabel('假正率 (FPR)', fontsize=12)
    plt.ylabel('真正率 (TPR)', fontsize=12)
    plt.title('特征曲线 (ROC)', fontsize=14, fontweight='bold')

    # 美化图例和添加文本注释
    plt.legend(loc="lower right", fontsize=10)
    plt.text(0.6, 0.3,
             f'最佳阈值: {optimal_threshold:.3f}\n对应TPR: {tpr[optimal_idx]:.3f}\n对应FPR: {fpr[optimal_idx]:.3f}',
             bbox=dict(facecolor='white', alpha=0.8))

    # 设置坐标轴刻度和样式
    plt.xticks(fontsize=10)
    plt.yticks(fontsize=10)
    plt.tight_layout()  # 确保所有元素都适合图表区域

    # 保存高清图片
    plt.savefig('../data/fig/ROC_curve_test.png', dpi=300, bbox_inches='tight')
    plt.show()
if __name__ == '__main__':
    pm = PowerLoadModel('../data/train.csv')
    feature_data, feature_columns = feature_engineering(pm.data_source)
    model_test('../data/test2.csv', feature_columns)